"""
模型定义模块
"""

import torch
import torch.nn as nn
import torchvision.models as models
from torchvision.models.resnet import ResNet50_Weights
from config import Config

class SuperclassClassifier(nn.Module):
    """
    CIFAR100超类分类模型
    基于预训练的ResNet50，替换最后的全连接层
    """
    def __init__(self):
        super(SuperclassClassifier, self).__init__()
        cfg = Config()
        
        # 加载预训练的ResNet50
        self.base_model = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
        
        # # 冻结前面的层(可选)
        # for param in self.base_model.parameters():
        #     param.requires_grad = False
        
        # 替换最后的全连接层
        num_ftrs = self.base_model.fc.in_features
        self.base_model.fc = nn.Sequential(
            nn.Linear(num_ftrs, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, cfg.NUM_CLASSES)
        )
        
    def forward(self, x):
        return self.base_model(x)
    
    def unfreeze_layers(self, num_layers=5):
        """
        解冻最后几层以便微调
        """
        children = list(self.base_model.children())
        for child in children[-num_layers:]:
            for param in child.parameters():
                param.requires_grad = True

def get_model():
    """
    获取模型实例
    """
    cfg = Config()
    model = SuperclassClassifier().to(cfg.DEVICE)
    return model